package com.tang.process;

import com.tang.bean.WaterSensor;
import com.tang.functions.WaterSensorMapFunction;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 实现TopN思路2, KeyedProcessFunction实现TopN
 * ==============
 * 1、按照vc做keyby，开窗，分别count
 *  => 增量聚合，计算 count
 *  => 全窗口，对计算结果 count值封装 ，  带上 窗口结束时间的 标签
 *  => 为了让同一个窗口时间范围的计算结果到一起去
 * 2、对同一个窗口范围的count值进行处理： 排序、取前N个
 *  => 按照 windowEnd做keyby
 *  => 使用process， 来一条调用一次，需要先存，分开存，用HashMap,key=windowEnd,value=List
 *  => 使用定时器，对 存起来的结果 进行 排序、取前N个
 *
 * @author tang
 * @since 2023/07/07 17:22
 */
public class KeyedProcessFunctionTopNDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(1);

        SingleOutputStreamOperator<WaterSensor> sensorDataStream = env.socketTextStream("192.168.70.141", 7777)
                .map(new WaterSensorMapFunction())
                .assignTimestampsAndWatermarks(
                        // 乱序的水位线，等待3s
                        WatermarkStrategy.<WaterSensor>forBoundedOutOfOrderness(Duration.ofSeconds(3))
                                // 指定时间戳分配器，从数据中提取，时间戳是 传感器记录时间戳TS * 1000，因为返回的时间戳是毫秒
                                .withTimestampAssigner(((element, recordTimestamp) -> element.getTs() * 1000L))
                );

        // 按照VC来keyBy
        SingleOutputStreamOperator<Tuple3<Integer, Integer, Long>> aggregate = sensorDataStream.keyBy(WaterSensor::getVc)
                // 开窗 窗口为10ms，滑动步长为5L
                .window(SlidingEventTimeWindows.of(Time.seconds(10L), Time.seconds(5L)))
                // 聚合 aggregate（增量聚合，全窗口打标签）
                .aggregate(new VcCountAgg(), new WindowResult());

        // 2. 按窗口标签（窗口结束时间）进行keyBy，保证同一个窗口时间范围的结果，到一起去。排序、取TopN
        aggregate.keyBy(r -> r.f2)
                .process(new TopN(2))
                .print();

        env.execute();
    }

    public static class VcCountAgg implements AggregateFunction<WaterSensor, Integer, Integer> {

        @Override
        public Integer createAccumulator() {
            return 0;
        }

        @Override
        public Integer add(WaterSensor value, Integer accumulator) {
            return accumulator + 1;
        }

        @Override
        public Integer getResult(Integer accumulator) {
            return accumulator;
        }

        @Override
        public Integer merge(Integer a, Integer b) {
            return null;
        }

    }

    public static class WindowResult extends ProcessWindowFunction<Integer, Tuple3<Integer, Integer, Long>, Integer, TimeWindow> {

        /**
         * 这里的数据是聚合后的数据在这里来哦
         *
         * @param key      键
         * @param context  上下文
         * @param elements 在窗口中被评估的元素
         * @param out      输出
         */
        @Override
        public void process(Integer key,
                            ProcessWindowFunction<Integer, Tuple3<Integer, Integer, Long>, Integer, TimeWindow>.Context context,
                            Iterable<Integer> elements, Collector<Tuple3<Integer, Integer, Long>> out) {
            // 迭代器里面只有1条数据，next一次即可
            System.out.println("触发了处理函数");
            Integer count = elements.iterator().next();
            long end = context.window().getEnd();
            out.collect(Tuple3.of(key, count, end));
        }

    }

    public static class TopN extends KeyedProcessFunction<Long, Tuple3<Integer, Integer, Long>, String> {

        private final Map<Long, List<Tuple3<Integer, Integer, Long>>> dataListMap;
        private final int threshold;

        public TopN(int threshold) {
            this.threshold = threshold;
            dataListMap = new ConcurrentHashMap<>();
        }

        @Override
        public void processElement(Tuple3<Integer, Integer, Long> value,
                                   KeyedProcessFunction<Long, Tuple3<Integer, Integer, Long>, String>.Context ctx,
                                   Collector<String> out) {
            Long currentKey = ctx.getCurrentKey();
            if (dataListMap.containsKey(currentKey)) {
                List<Tuple3<Integer, Integer, Long>> list = dataListMap.get(currentKey);
                list.add(value);
            } else {
                List<Tuple3<Integer, Integer, Long>> list = new ArrayList<>();
                list.add(value);
                dataListMap.put(currentKey, list);
            }

            // 注册一个1ms后的事件时间定时器即可
            ctx.timerService().registerEventTimeTimer(currentKey + 1);

        }

        @Override
        public void onTimer(long timestamp,
                            KeyedProcessFunction<Long, Tuple3<Integer, Integer, Long>, String>.OnTimerContext ctx,
                            Collector<String> out) throws Exception {
            super.onTimer(timestamp, ctx, out);
            Long currentKey = ctx.getCurrentKey();

            // 1.排序
            List<Tuple3<Integer, Integer, Long>> list = dataListMap.get(currentKey);
            list.sort((o1, o2) -> o2.f1 - o1.f1); // 根据count排序
            // 3.取出count最大的两个vc
            StringBuilder outBuilder = new StringBuilder();

            outBuilder.append("=========================\n");
            // 3.1 遍历排序后的list，取出前两个，考虑可能list只有一个的情况，这里还是取个list size和2的最小。
            for (int i = 0, len = Math.min(threshold, list.size()); i < len; i++) {
                Tuple3<Integer, Integer, Long> vcTuple = list.get(i);
                outBuilder.append("Top").append(i + 1).append("\n")
                        .append("vc = ").append(vcTuple.f0).append("\n")
                        .append("count = ").append(vcTuple.f1).append("\n")
                        .append("窗口结束时间 = ").append(currentKey).append("\n")
                        .append("=========================\n");
            }

            // 用完List就清空，让gc生效
            list.clear();

            out.collect(outBuilder.toString());
        }

    }

}
